﻿using System.Reflection;
using Faml.Frontend;
using Faml.Operations;
using System.Collections.Generic;

namespace Faml.Optimization
{
    public class GetLengthSubarrayVisitor : IVisitor
    {
        #region IVisitor Members

        public FamlExpression Visit(FamlExpression node)
        {
            if (node.NodeType != FamlExpressionType.FunctionCall)
                return node;

            FamlFunctionCallExpression mce = (FamlFunctionCallExpression)node;
            string opAName = mce.Function;

            if (string.Compare(opAName, "length", true) != 0)
                return node;

            FamlExpression arg0 = mce.Arguments[0];
            if (arg0.NodeType != FamlExpressionType.FunctionCall)
                return node;

            FamlFunctionCallExpression sub = arg0 as FamlFunctionCallExpression;
            string opBName = sub.Function;
            if (string.Compare(opBName, "subarray", true) != 0)
                return node;

            FamlExpression index_array = null;

            // length(subarray(A, index_array))  ===> length(index_array)
            if (mce.Arguments.Count == 1)
            {
                index_array = FindIndexArray(sub.Arguments, 0);
            }
            // length(subarray(A, ...index_array...), dim)  ===> length(index_array)
            else if (mce.Arguments.Count == 2)
            {
                FamlExpression ex_dim = mce.Arguments[1];
                if (ex_dim.NodeType != FamlExpressionType.Constant)
                    return node;

                FamlConstant ce = ex_dim as FamlConstant;
                if (ce.Type != typeof(int))
                    return node;
                int dim = (int)ce.Value;

                index_array = FindIndexArray(sub.Arguments, dim);
                if (index_array == null)
                    return node;
            }            

            if (index_array == null)
                return node;

            MethodInfo[] getLen = typeof(GenericGetLength).GetMethods();
            for (int i = 0; i < getLen.Length; i++)
            {
                if (getLen[i].GetParameters().Length == 1)
                {
                    MethodInfo m = getLen[i].MakeGenericMethod(typeof(int));
                    return new FamlFunctionCallExpression(new Operation("length", m), index_array);
                }
            }

            return node;
        }

        #endregion

        private FamlExpression FindIndexArray(IList<FamlExpression> args, int n)
        {
            bool first = true;
            foreach (FamlExpression arg in args)
            {
                if (first)
                {
                    first = false;
                    continue;
                }
                if (arg.Type == typeof(int[]))
                {
                    if (n-- == 0)
                    {
                        return arg;
                    }
                }
            }
            return null; 
        }
    }
}
